import torch
from thop import profile
from thop import clever_format
import argparse

from models.experimental import attempt_load
from utils.torch_utils import select_device

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='/home/lzc/work/Weights/yolov7-off/v0.1/yolov7.pt', help='weights path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()

    device = select_device(opt.device)
    model = attempt_load(opt.weights, map_location=device)  # load FP32 model

    input = torch.randn(1, 3, opt.img_size[0], opt.img_size[1])
    flops, params = profile(model, inputs=(input,), verbose=False)
    print(flops, params)
    flops, params = clever_format([flops, params], "%.3f")
    print(flops, params)
